/** topology.c - cs613011--Harry Ticer*/
#include <stdlib.h>
#include <apr_pools.h>
#include "mpi.h"
#include "env.h"
#include "oom.h"
#include "cm.h"
#include "log.h"


/** structure to hold cartesian info*/
typedef struct {
	int ndims;
	int *dims;
	int *periods;
	int reorder;
	int own_position;
} Cartinfo;

/**
 *	We need a constant key for the hashing of the cartinfo structure 
 *      to the cart comm object.
 *	Going to use a static string variable.
 */
static char *cartinfo_key = "cartinfo";

/** ----ignore this struct for now---- */
struct TPG {
	/* total processes....also size of the array */
	int nranks;
	/* these next two will have memory allocated via malloc(), based on nranks
		This will happen after a special function call outside of the MPI spec	*/
	int *rank;
	float *cost;
};

/** Makes a new communicator to which topology information has been 
 *attached, and currently ignores reorder (assumes false) */

/* TODO: make the following functions work with periodic flag....and check 
	array bounds (right now, the following functions assume that the numbers given will 
	always be in bounds)
*/

int MPI_Cart_create (MPI_Comm comm_old, int ndims, int *dims, int *periods, int reorder, MPI_Comm *comm_cart) {
	/** Notes:
	 *	Not going to actually have an N-dimensional array. Will 
         *      use just 1 dimension, and
	 *	access it based on what it should be by using location calculations.
	 */

	/** returns error if ndims doesn't make sense **/
	if (ndims <= 0) {
		return MPI_ERR_DIMS;
	}

	/*	other init	*/
	int i;
	int totaldim = 1;
	int comm_old_size;
	int *ranklist;
	MPI_Group comm_old_group;
	MPI_Group cart_group;
	Cartinfo *cartinfo;
	Group *realgroup_old;
	Comm *realcart_comm;

	apr_pool_t *mempool;
	apr_pool_create(&mempool, NULL);

	MPI_Comm_size(comm_old, &comm_old_size);
	MPI_Comm_group(comm_old, &comm_old_group);

	realgroup_old = OOM_group_lookup(comm_old_group);

	/** set up Cartinfo struct stuff */
	cartinfo = OOM_palloc(sizeof(Cartinfo));
	cartinfo->ndims = ndims;
	cartinfo->dims = OOM_palloc(ndims * sizeof(int));
	cartinfo->periods = OOM_palloc(ndims * sizeof(int));
	cartinfo->reorder = reorder;
	for (i = 0; i < ndims; i++) {
		cartinfo->dims[i] = dims[i];
		cartinfo->periods[i] = periods[i];
	}

	/** ignore reorder....*/
	MPI_Comm_rank(comm_old, &(cartinfo->own_position));

	/** count up total processes required */
	for (i = 0; i < ndims; i++) {
		totaldim *= dims[i];
	}

	/** set up an array for the ranks of the possible new group/communicator */
	ranklist = (int *) apr_palloc(mempool, totaldim * sizeof(int));

	/** if grid size is bigger than comm_old group size, return error */
	if (totaldim > comm_old_size) {apr_pool_destroy(mempool); return MPI_ERR_DIMS;}
	else {
		/** copy ranks from old_comm's group. doing this after 
         *  checking if totaldim > comm_old_size,
		 *  so that we ensure that we don't read out of bounds on the old comm's group
		 */
		for (i = 0; i < totaldim; i++) {
			ranklist[i] = realgroup_old->members[i];
		}
		/** just duplicate communicator if they're the same size  */
		if (totaldim == comm_old_size) {
			MPI_Comm_dup(comm_old, comm_cart);
		}
		else {
 		/** if grid size is smaller than comm_old group size, some processes 
  		*  return MPI_COMM_NULL
  		*/
		
		/** find my rank, see I'm part of the new group, return null if not, 
		 *	assign rest to new group
		 */
		/*for (i = 0; i < totaldim; i++) {
			if (ranklist[i] == cartinfo->own_position) {break;}
		}
		if (i >= totaldim) {apr_pool_destroy(mempool); return MPI_COMM_NULL;}
		*/
		/* communicator won't exist for the process that called it if its rank
			is not included in the group
		*/
		MPI_Group_incl(comm_old_group, totaldim, ranklist, &cart_group);
		MPI_Comm_create(comm_old, cart_group, comm_cart);

		}
	}

	apr_pool_destroy(mempool);

	if (*comm_cart != MPI_COMM_NULL)
	{
		/** set the topology variable of the new communicator to MPI_CART */
		realcart_comm = OOM_comm_lookup(*comm_cart);
		realcart_comm->topology = MPI_CART;

	/** now attach the cartinfo struct to comm_cart...need handle mechanism for 
	 *	attributes hash table
	 */

		apr_hash_set(realcart_comm->attributes, cartinfo_key, APR_HASH_KEY_STRING, cartinfo);
	}

	return MPI_SUCCESS;
}

/** Retrieves Cartesian topology information associated with a communicator */
int MPI_Cart_get (MPI_Comm comm, int maxdims, int *dims, int *periods, int *coords) {

	/* TODO: check array size errors */

	int i;
	Cartinfo *cartinfo;
	/** check to see if it actually is a cartesian communicator */
	Comm *realcart_comm = OOM_comm_lookup(comm);
	if (realcart_comm->topology != MPI_CART) {return MPI_ERR_TOPOLOGY;}

	/** retrieve attached Cartinfo */
	cartinfo = (Cartinfo*) apr_hash_get(realcart_comm->attributes, cartinfo_key, APR_HASH_KEY_STRING);

	for (i = 0; i < maxdims; i++) {
		dims[i] = cartinfo->dims[i];
		periods[i] = cartinfo->periods[i];
		MPI_Cart_coords (comm, cartinfo->own_position, maxdims, coords);
	}

	return MPI_SUCCESS;
}

/**-----MPI_...Cart_coords, Cart_shift do not function right now----- */
/* TODO: make the following functions work with periodic flag */
/* Retrieves dimm information associated with a communicator */
int MPI_Cartdim_get (MPI_Comm comm, int *ndims) {
/* Determines process rank in communicator given Cartesian location...only works for 2d grid at the moment
		todo: code a more general algorithm
	/* TODO: check array size errors */

	int i;
	Cartinfo *cartinfo;
	/* check to see if it actually is a cartesian communicator */
	Comm *realcart_comm = OOM_comm_lookup(comm);
	if (realcart_comm->topology != MPI_CART) {return MPI_ERR_TOPOLOGY;}

	/* retrieve attached Cartinfo */
	cartinfo = (Cartinfo*) apr_hash_get(realcart_comm->attributes, cartinfo_key, APR_HASH_KEY_STRING);

	*ndims = cartinfo->ndims;

	return MPI_SUCCESS;
}

/*-----Cart_shift does not function right now----- */

/** Determines process rank in communicator given Cartesian location */
int MPI_Cart_rank (MPI_Comm comm, int *coords, int *rank) {

	int i;
	int j;
	/*int totaldim = 1;*/
	int thisdim = 1;
	int grandtotal = 0;
	Cartinfo *cartinfo;

	MPI_Group fakegroup;
	Group *realgroup;

	/** check to see if it actually is a cartesian communicator */
	Comm *realcart_comm = OOM_comm_lookup(comm);
	if (realcart_comm->topology != MPI_CART) {return MPI_ERR_TOPOLOGY;}

	MPI_Comm_group(comm, &fakegroup);
	realgroup = OOM_group_lookup(fakegroup);

	/** retrieve attached Cartinfo */
	cartinfo = (Cartinfo*) apr_hash_get(realcart_comm->attributes, cartinfo_key, APR_HASH_KEY_STRING);

	/** calculate rank */
	int cndims = cartinfo->ndims;
	
	for (i = 0; i < cndims-1; i++) {
		for (j = cndims - 1; j > 0 + i; j--) {
			thisdim *= cartinfo->dims[j];
		}
		thisdim *= coords[0 + i];
		grandtotal += thisdim;
		thisdim = 1;
	}
	grandtotal += coords[cndims - 1];
	
	if ((grandtotal+1) > realgroup->num_members) {return MPI_ERR_RANK;}
	*rank = realgroup->members[grandtotal];
	
	return MPI_SUCCESS;
}
/*	x, y, z

multiply all dimensions, their max coords, times the deepest dimension

get that number, add it to the above, with the deepest as one dimension less

loop until you get to the smallest dimension, and then just add that coord
*/

/** Determines process coords in cartesian topology given rank in group 
 */
int MPI_Cart_coords (MPI_Comm comm, int rank, int maxdims, int *coords) {

	MPI_Group fakegroup;
	Group *realgroup;
	Cartinfo *cartinfo;
	int groupsize;
	int rankloc;
	int i, j;
	int FinalIndex;
	int thisdimloc;
	int *sizes;
	int *reversed;
	
	apr_pool_t *mempool;
	apr_pool_create(&mempool, NULL);
	
	sizes = (int *) apr_palloc(mempool, maxdims * sizeof(int));
	reversed = (int *) apr_palloc(mempool, maxdims * sizeof(int));

	/* check to see if it actually is a cartesian communicator */
	Comm *realcart_comm = OOM_comm_lookup(comm);
	if (realcart_comm->topology != MPI_CART) {return MPI_ERR_TOPOLOGY;}

	MPI_Comm_group(comm, &fakegroup);
	realgroup = OOM_group_lookup(fakegroup);
	cartinfo = (Cartinfo*) apr_hash_get(realcart_comm->attributes, cartinfo_key, APR_HASH_KEY_STRING);
	MPI_Comm_size(comm, &groupsize);

	/* check to see if maxdims makes any sense */
	if (maxdims != cartinfo->ndims) {return MPI_ERR_DIMS;}

	for (rankloc = 0; rankloc < groupsize; rankloc++) {
		if (rank == realgroup->members[rankloc]) {break;}
	}
	if (rankloc == groupsize) {
		/* rank not in group */
		return MPI_ERR_RANK;
	}

	/* figure out CSize, BSize, etc... */
	for (i = 0; i < maxdims-1; i++) {
		sizes[i] = 1;
		for (j = 1+i; j < maxdims; j++) {
			sizes[i] *= cartinfo->dims[j];
		}
	}	

	FinalIndex = 0;	
	for (i = 0; i < maxdims-1; i++) {			
		thisdimloc = 1;	
		while (1) {
			if ((thisdimloc * sizes[i] + FinalIndex) > rankloc) {
				thisdimloc--;
				FinalIndex += thisdimloc * sizes[i];
				coords[i] = thisdimloc;
				break;
			}
			else {
				thisdimloc++;
			}
		}
	}
	coords[maxdims-1] = rankloc-FinalIndex;

	apr_pool_destroy(mempool);
	return MPI_SUCCESS;
}

/** Returns the shifted source and destination ranks, given a shift 
direction 
and amount */
int MPI_Cart_shift (MPI_Comm comm, int direction, int displ, int *source, int *dest) {

	return 0;
}

/* For cartesian topologies, the function MPI_DIMS_CREATE helps
	the user select a balanced distribution of processes per coordinate
	direction, depending on the number of processes in the group to be
	balanced and optional constraints that can be specified by the user.
	One use is to partition all the processes (the size of MPI_COMM_WORLD's
	group) into an n-dimensional topology. 
*/
/* this function works only on a very basic level right now */
int MPI_Dims_create(int nnodes, int ndims, int *dims) {
	
	int i;

	if (ndims < 1) {return MPI_ERR_DIMS;}

	if (nnodes % ndims == 0 && ndims != 1) {
		if (nnodes/ndims > ndims) {
			dims[0] = nnodes/ndims;
			dims[1] = ndims;
		}
		for (i = 2; i < ndims; i++) {
			dims[i] = 1;
		}
	}
	else {
		dims[0] = nnodes;
	}

	return MPI_SUCCESS;
	
}
